import pygltflib as gltf2
import numpy as np

# dirty patch for numpy high version
np.bool = np.bool_

from vtkmodules.vtkCommonDataModel import (
    vtkAnimationScene,
)


from vtkmodules.vtkIOImage import (
    vtkJPEGReader,
    vtkPNGReader
)

from vtkmodules.vtkCommonTransforms import vtkTransform
from vtkmodules.vtkImagingCore import vtkImageFlip
from vtkmodules.vtkCommonMath  import vtkMatrix4x4
from vtkmodules.vtkRenderingOpenGL2 import (
    vtkOpenGLTexture,
    vtkTextureObject
) 

from .Mesh import *
from .Material import *
from .Node import *
from .Animation import *
from .Scene import *
from .Utils import *

class GLTFLoader:

    def __init__(self):
        self.gltf = None

    def Load(self, path):
        self.gltf = gltf2.GLTF2().load(path)
        #self.gltf.convert_buffers(gltf2.BufferFormat.DATAURI) 
        
        # load images
        self.images = []
        for image in self.gltf.images:            
            self.images.append(self._parse_image(image))

        # load textures
        self.textures = []        
        for tex in self.gltf.textures:
            self.textures.append(self._parse_texture(tex))

        self.materials = []
        for mtl in self.gltf.materials:
            self.materials.append(Material(mtl, self.textures))

        # parse meshes
        self.meshes = []
        for mesh in self.gltf.meshes:
            self.meshes.append(self._parse_mesh(mesh))

        # parse nodes
        self.nodes = []
        for node in self.gltf.nodes:
            self.nodes.append(self._parse_node(node))

        # second pass, build node relation
        for idx in range(len(self.gltf.nodes)):
            node = self.nodes[idx]   
            for childIdx in self.gltf.nodes[idx].children:
                node.AddChild(self.nodes[childIdx])

        # parse animations
        self.animations = []
        for animation in self.gltf.animations:          
            self.animations.append(self._parse_animation(animation))             
        
        # parse skins
        self.skins = []
        for skin in self.gltf.skins:            
            self.skins.append(self._parse_skin(skin))

        # update skin node
        for idx in range(len(self.gltf.nodes)):
            node = self.nodes[idx] 
            id = self.gltf.nodes[idx].skin
            if id is not None:
                node.skin = self.skins[id]

        self._update_skin_animation()
        
        print("Load GLTF File:", path)
        print(" #Scene:   ", len(self.gltf.scenes))
        print(" #Node:    ", len(self.nodes))
        print(" #Mesh:    ", len(self.meshes))
        print(" #Material:", len(self.materials))

        # count vertices and triangles
        polys = 0
        verts = 0
        for node in self.nodes:
            for mesh in node.meshes:
                for prim in mesh.primitives:
                    verts += prim.polyData.GetPoints().GetNumberOfPoints()
                    polys += prim.polyData.GetNumberOfPolys()        
        print(" #Vertex:  ", verts)
        print(" #Triangle:", polys)
        print(" #Animation:", len(self.animations))


    def _parse_mesh(self, m):
        mesh = Mesh()
        mesh.Parse(self.gltf, m)
        # mesh.name = m.name

        # # mesh may contain multiple primitives...
        # for p in m.primitives:
        #     prim = Primitive()
        #     prim.ParseGeometry(self.gltf, p) 
        #     # setup materials
        #     prim.material = self.materials[p.material]
        #     mesh.AddPrimitive(prim)

        return mesh


    def _parse_image(self, img):
        if img.mimeType == 'image/jpeg':
            reader = vtkJPEGReader()
        elif img.mimeType == 'image/png':
            reader = vtkPNGReader()

        bv = self.gltf.bufferViews[img.bufferView]
        bf = self.gltf.get_data_from_buffer_uri(self.gltf.buffers[bv.buffer].uri)
        data = bf[bv.byteOffset: bv.byteOffset + bv.byteLength]

        reader.SetMemoryBuffer(data)
        reader.SetMemoryBufferLength(bv.byteLength)
        reader.Update()

        # flip y axis...
        flip = vtkImageFlip()
        flip.SetFilteredAxis(1)
        flip.SetInputConnection(reader.GetOutputPort())
        flip.Update()
        
        image = flip.GetOutput()
        image.name = img.name
        return image

    def _parse_texture(self, tex):
    
        texture = vtkOpenGLTexture()
        texture.name = tex.name

        texObj = vtkTextureObject() 
        # TODO: parse sampler
        sampler = self.gltf.samplers[tex.sampler]  
        # magnifacation filters
        if sampler.magFilter == gltf2.NEAREST:
            texObj.SetMagnificationFilter(vtkTextureObject.Nearest)
        elif sampler.magFilter == gltf2.LINEAR:
            texObj.SetMagnificationFilter(vtkTextureObject.Linear)

        # nimification filters
        if sampler.minFilter == gltf2.NEAREST:            
            texObj.SetMinificationFilter(vtkTextureObject.Nearest)
        elif sampler.minFilter == gltf2.LINEAR:
            texObj.SetMinificationFilter(vtkTextureObject.Linear)
        elif sampler.minFilter == gltf2.NEAREST_MIPMAP_NEAREST:
            texObj.SetMinificationFilter(vtkTextureObject.NearestMipmapNearest)
        elif sampler.minFilter == gltf2.NEAREST_MIPMAP_LINEAR:
            texObj.SetMinificationFilter(vtkTextureObject.NearestMipmapLinear)
        elif sampler.minFilter == gltf2.LINEAR_MIPMAP_NEAREST:
            texObj.SetMinificationFilter(vtkTextureObject.LinearMipmapNearest)
        elif sampler.minFilter == gltf2.LINEAR_MIPMAP_LINEAR:
            texObj.SetMinificationFilter(vtkTextureObject.LinearMipmapLinear)

        if sampler.wrapS == gltf2.CLAMP_TO_EDGE:
            texObj.SetWrapS(vtkTextureObject.ClampToEdge)
        elif sampler.wrapS == gltf2.MIRRORED_REPEAT:
            texObj.SetWrapS(vtkTextureObject.MirroredRepeat)
        elif sampler.wrapS == gltf2.REPEAT:
            texObj.SetWrapS(vtkTextureObject.Repeat)

        if sampler.wrapT == gltf2.CLAMP_TO_EDGE:
            texObj.SetWrapT(vtkTextureObject.ClampToEdge)
        elif sampler.wrapT == gltf2.MIRRORED_REPEAT:
            texObj.SetWrapT(vtkTextureObject.MirroredRepeat)
        elif sampler.wrapT == gltf2.REPEAT:
            texObj.SetWrapT(vtkTextureObject.Repeat)
        # TODO: set texture object for sampler

        # always generate mipmap...
        texture.MipmapOn() 
        texture.InterpolateOn()
        # set mipmap and interoplation before load data
        texture.SetInputData(self.images[tex.source]) 
        return texture

    def _parse_node(self, n):        
        node = Node()
        node.name = n.name
        node.extras = n.extras

        if n.matrix != None:
            # col major from gltf to row major in vtk
            tm = [0.0] * 16
            vtkMatrix4x4.Transpose(n.matrix, tm)   
            node.transform.SetMatrix(tm)
        else:
            if n.translation != None:
                T = n.translation  
                node.translate.Translate(T)
            if n.scale != None:
                S = n.scale
                node.scale.Scale(S)
            if n.rotation != None:
                node.rotate.SetMatrix(get_rotation_mat(n.rotation)) 

        if n.mesh != None:
            node.AddMesh(self.meshes[n.mesh])
            
        return node

    def _parse_animation(self, gltfAnimation):
        # use vtkAnimationScene and vtkAnimationCue
        animation = Animation(gltfAnimation.name)
        
        start = 1000    # hack...
        end   = 0
        for channel in gltfAnimation.channels:

            sampler = gltfAnimation.samplers[channel.sampler]

            ac = AnimationChannel(
                self.nodes[channel.target.node],
                channel.target.path,
                get_accessor_data(self.gltf, sampler.input),
                get_accessor_data(self.gltf, sampler.output),
                sampler.interpolation
            )
            animation.AddCue(ac)

            start = min(start, ac.GetStartTime())
            end   = max(end,   ac.GetEndTime())

        animation.SetStartTime(start)
        animation.SetEndTime(end)        
        animation.AddObserver("AnimationCueTickEvent", self._update_skin_animation)
        return animation

    def _update_skin_animation(self, info=None, event=None):  
        for node in self.nodes:
            if node.skin == None:
                continue            

            # for those nodes with skin, skip the its own transform
            #node.worldTransform.Identity() 
            node.Identity() 

            for mesh in node.meshes:
                for prim in mesh.primitives:
                    prim.UpdateSkinTransform(node.skin['transforms'])


    def _parse_skin(self, skin):
        # get inverse bind matrices for vtk
        inverseBindMatrices = []
        invMatsData = get_accessor_data(self.gltf, skin.inverseBindMatrices)

        skinTransforms = []
        
        for i in range(len(skin.joints)):
            invMat = vtkTransform()
            invMat.SetMatrix(np.transpose(invMatsData[i]).flatten())
            inverseBindMatrices.append(invMat)

            jointMat = self.nodes[skin.joints[i]]#.worldTransform
            tm = vtkTransform()
            tm.Concatenate(jointMat)
            tm.Concatenate(invMat)
            skinTransforms.append(tm)

        return {
            "skeleton" : self.nodes[skin.skeleton],
            "joints": [self.nodes[id] for id in skin.joints],
            "inverseBindMatrices": inverseBindMatrices,
            "transforms" : skinTransforms
        }

    
    def GetScene(self, idx=None):
        # default scene
        if idx is None:
            idx = self.gltf.scene

        # first scene...
        if idx is None:
            idx = 0

        if idx >= 0 and idx < len(self.gltf.scenes):
            scn = self.gltf.scenes[idx]
        else:
            return None
        
        # create scene
        scene = Scene()
        scene.materials = self.materials
        scene.animations = self.animations
        
        for idx in scn.nodes:
            scene.AddNode(self.nodes[idx])

        scene.Update()
        return scene
        

